import os  
import torch
import torchvision
from utils.utils import get_accurary
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader 
import numpy as np
from torchvision import transforms
from PIL import Image  
import seaborn as sns
from dataset.MultiLabelImageDatasetTest import MultiLabelImageDatasetTest
from datetime import datetime, timedelta
from utils.utils import set_seed
from config import opt
from scipy import spatial

res = ""
set_seed(opt.seed)

# 转换为Tensor数据并归一化
# 定义预处理步骤
transformer = transforms.Compose([
    
    transforms.Resize(256),
    transforms.CenterCrop(224),  # 通常在调整大小后进行中心裁剪

    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])




net = torch.load(opt.model_path)
net = net.to(opt.device)
net.eval()

# 数据加载器,DataLoader将输入的数据按照batch_size封装成tensor
# 然后，后续需要再包装成Variable就可以作为模型的输入 了
test_dataset = MultiLabelImageDatasetTest(opt.path + 'test/', transform=transformer)

test_loader = DataLoader(test_dataset, batch_size=opt.test_batch_size, shuffle=False)

distance = 0.615
print("test")
known_dis_list = []
all_dis = []
unkown_dis_list = []
acc = 0
acc1 = 0
acc2 = 0
acc3 = 0
num = 0
acc4 = 0
threshold = 0.5
with torch.no_grad():

    for i, data_info in enumerate(test_loader):

        images, labels = data_info['image_tensors'].to(opt.device), data_info['labels']
        outputs,trans_feature,features = net(images)

        
        keys_to_select1 = opt.known_sig
        selected_outputs1 = {k: outputs[k] for k in keys_to_select1}
        selected_labels1 = {k: labels[k] for k in keys_to_select1}
        acc4 += get_accurary(selected_outputs1, selected_labels1, "cuda",threshold)

        outputs['T000']  = torch.zeros_like(outputs['T001'])
        acc1 += get_accurary(outputs, labels, "cuda",threshold)
        features = features.detach().cpu().numpy()

        ave_trans_feature = np.zeros_like(features)
        for i in trans_feature.keys():
            ave_trans_feature  += trans_feature[i][0].detach().cpu().numpy()
        ave_trans_feature = ave_trans_feature/3


        dis =  1 - spatial.distance.cosine(ave_trans_feature[0] , features[0])

        if(labels['T000'].item() == 1):
            unkown_dis_list.append(dis)
        else:
            known_dis_list.append(dis)

        all_dis.append(dis)
        
        keys_to_select = ['T000']
        selected_outputs = {k: outputs[k] for k in keys_to_select}
        selected_labels = {k: labels[k] for k in keys_to_select}

        acc += get_accurary(selected_outputs, selected_labels, "cuda",threshold)
        # print(dis)
        if dis >distance:
            outputs['T000'] = torch.tensor([[1.]]).cuda()
        else:
            outputs['T000'] = torch.tensor([[0.]]).cuda()

        selected_outputs = {k: outputs[k] for k in keys_to_select}
        # print(selected_outputs)
        acc2 += get_accurary(selected_outputs, selected_labels, "cuda",threshold)
        acc3 += get_accurary(outputs, labels, "cuda",threshold)
        num +=1
# acc/num
print("全0时未知信号acc,全0时所有信号acc1,阈值判断后的未知信号acc2,阈值判断后的所有信号acc3,已知信号acc4")
print( acc/num,acc1/num,acc2/num,acc3/num,acc4/num)
res += "test\n全0时未知信号acc,全0时所有信号acc1,阈值判断后的未知信号acc2,阈值判断后的所有信号acc3,已知信号acc4\n"
res += "%f,%f,%f,%f,%f\n"%(acc/num,acc1/num,acc2/num,acc3/num,acc4/num)




plt.hist(unkown_dis_list, bins=50, color="red", alpha=.9)
plt.hist(known_dis_list, bins=50,color="blue", alpha=.5)
plt.savefig("./picture/dis" + opt.log_time +"_res.png") # 你给图片起的名字

with open("./log/test/" + opt.log_time  +"_res.txt", "a", encoding="utf-8") as file:
    file.write(res)